#!/usr/bin/env python
""" 
Program to find all taxa containing a name of interest. For
example, if you want to find all nodes with Mycobacteria
in the name, this program will return the path to those
nodes on one line, and the ranks of each of those taxa
on the next.

Example:-
python {0} -in Mycobacteria Escherichia

""".format( __file__ )

from ssummolib import dict_walk,getRanks, TaxDB,load_index,Options
from sys import argv
from os.path import sep
import re


def find_taxa( tdict, options ):
    found = []
    interesting = re.compile ( r'|'.join( options['-in'] ),re.I )
    for path,node in dict_walk( '', tdict ):
        if interesting.search( path.rsplit( sep , 1 )[-1] ):
            found.append( path )
    return found

def match_ranks( taxa ):
    TaxDBObj = TaxDB()
    matchedRanks = {}
    for taxon in taxa:
        ranks,table = getRanks( taxon.split( sep ), TaxDBObj )
        matchedRanks.update( { taxon : ranks } )
    return matchedRanks

def printMatches( nameRanks ):
    taxDict = {'rank':'root'}
    for taxPath in nameRanks:
        node = taxDict
        taxPathList = taxPath.split( sep )
        for i in xrange(len( taxPathList )):
            taxon = taxPathList[i]
            rank = nameRanks[taxPath][i]
            if taxon in node:
                node = node[taxon]
            else:
                node.update( { taxon: {'rank': rank } } )
                node = node[taxon]
    answer = raw_input( 'We have {0} names matching your search criteria. Do you want to see them all??'.format( len(nameRanks) ) )
    if answer.startswith('y'):
        rankList = []
        for names,node,path in walk_ranks(taxDict):
            rank = node.pop( 'rank' )
            pathList = path.split( sep )
            depth = len( pathList )
            nRanks = len(rankList)
            if nRanks < depth:  ### Update ranks
                rankList.append(set( [rank] ))
            elif nRanks ==depth:
                pass
                #if rank not in rankList[depth-1]:
                #    rankList[depth-1].add(rank)
            elif rank not in rankList[depth]:
                rankList[depth].add( rank ) 
            else:
                pass
            nKids = len(names)
            if nKids == 1:
                print '  ' * depth,'|_', pathList[-1]
            elif nKids == 0:
                print '  ' * depth,'|-',pathList[-1]
            else:  # Multiple children
                print '  ' * depth,'|-',pathList[-1]
                continue
                ind = 0
                while ind < nKids-1:
                    print '  ' * depth,'|-',names[ind]
                    ind += 1
                print '  '* depth,'|_',names[-1]
        for i in xrange(1,len(rankList)):
            if len(rankList[i]) == 1:
                print '  '*i,'|-',rankList[i].pop()
            else:   # multiple
                ranks = []
                for rank in rankList[i]:
                    ranks.append(rank )
                print '  '*i,'|-', ', '.join( ranks )

    else:
        mainPath = ['','']
        for names,node,path in walk_ranks( taxDict ):
            if len(names) == 1:
                mainPath[0] += names[0]
                mainPath[1] += node['rank']
            else:
                break
        print 'The single undiverging branch we went down looked like:-'
        print '\n'.join( mainPath )
        print '\nThen we had the choice of:-'
        print names
        print node['rank']
            

def walk_ranks( tdict,path='' ):
    names = tdict.keys()
    del( names[names.index('rank')] )
    yield names,tdict,path
    for name in names:
        for x in walk_ranks(tdict[name], path=path+sep+name ):
            yield x


def main(args):
    tdict = load_index()
    options = parse_args( args )
    taxaNames = find_taxa( tdict, options )
    nameRankDict = match_ranks( taxaNames )
    printMatches( nameRankDict )

def parse_args( args ):
    options = Options()
    options.multiargs = ['-in']
    options.singleargs= []
    options.helpTxt = {'-in': 'Taxa name(s) of interest'}
    options.parseArgs( argv[1:] )
    return options

if __name__ == '__main__':
    main(argv[1:])

